from tqdm import tqdm
from yan import CvFo
import paddle
import pandas as pd
import numpy as np
def gen_data(data,voc_l,pos_list):
    data_list=[]
    data_len=[]

    for i in data:
        i=i.split()
        i=["<|aos|>"]+i+["<|bos|>"]
        data_list.append(i)
        data_len.append(len(i))
    max_len=max(data_len)+1
    data=[]
    for i in data_list:
        i+= pos_list[len(i):max_len]
        i=[voc_l.get(j,0) for j in i]
        data.append(i)
    return data


voc = pd.read_pickle("E:/CVFO/yan/simplebooks_92/vocab.pkl")
voc_dict={ j:i for i,j in enumerate(voc)}
net = CvFo(len(voc), 16, 16)
loss_func = paddle.nn.CrossEntropyLoss()
loss_func1 = paddle.nn.MSELoss()
optimizer = paddle.optimizer.Adam(parameters=net.parameters(), learning_rate=0.0001)
with open('E:/CVFO/yan/simplebooks_92/train.txt', 'r', encoding="utf-8") as f:
    lines = f.readlines()
data_set = [i.strip() for i in tqdm(lines) if len(i) > 2 and len(i.strip().split()) < 136]
pos=["<|p_{}|>".format(i) for i in range(139)]

batch_size = 12
bar = tqdm(len(100 * list(range(0, len(data_set), batch_size))))
plt_loss=[]
for epoch in range(100):
    np.random.shuffle(data_set)

    loss_list=[]
    for i in range(0, len(data_set), batch_size):
        j = i + batch_size
        one_data = gen_data(data_set[i:j],voc_dict,pos)
        two_data = paddle.to_tensor(one_data).astype('int64')

        out= net(two_data[:, :-1])
        loss = loss_func(out.reshape([-1, out.shape[-1]]), two_data[:, 1:].reshape([-1]))

        loss_list.append(loss.item())

        if np.mean(loss_list)<1.6:
            plt_loss.append(np.mean(loss_list))
            if len(plt_loss)%10000==0:
                pd.to_pickle(plt_loss, "loss_star_star_star.pandas_pickle")

        bar.set_description("loss___{:.5f}".format(np.mean(loss_list)))
        bar.update(1)
        optimizer.clear_grad()
        loss.backward()
        optimizer.step()
    if (epoch + 1) % 10 == 0:
        paddle.save(net.state_dict(), "model_{}_loss_{:.5f}.pth".format(epoch + 1, np.mean(loss_list)))
paddle.save(net.state_dict(), "model_{}_loss_{:.5f}.pth".format(epoch + 1, loss.item()))
